#include <cstdio>
#include <ctime>
#include <cstring>
#include <cstdlib>
#include <iostream>
#include <vector>
#include <set>
#include <algorithm>
#include <iterator>

#include "types.h"
#include "Initial.h"
#include "Matrix.h"
#include "SimpleElementS.h"
#include "LinearEquationGroup.h"

const int sbox_size = 256;

/*byte sbox[256] = {
	141  ,   1  , 203  , 209  , 229  , 225  , 232  , 192  , 253  ,  63  , 255  , 178  ,  96  ,  43  , 116  ,  75  ,
	146  ,  69  , 243  ,  66  ,  89  , 187  , 242  , 111  , 152  ,  10  ,  48  , 194  , 168  ,  77  ,  58  , 241  ,
	175  ,  94  ,  73  ,  67  ,  33  , 147  , 244  , 223  , 208  , 112  , 161  , 130  , 186  , 181  , 121  , 133  ,
	5  ,  92  ,  76  , 191  ,  97  , 236  ,  24  , 240  , 171  , 100  ,  84  ,   9  , 245  ,  49  ,  29  , 103  ,
	47  ,  40  , 218  ,  15  , 172  , 252  , 169  ,   4  , 196  , 139  , 157  , 107  , 226  , 219  , 122  ,  99  ,
	56  ,  35  , 104  , 140  ,  65  ,  26  , 221  , 160  , 215  ,  78  ,  93  , 179  , 207  ,  14  , 177  , 235  ,
	46  , 124  , 143  ,  72  , 210  , 231  ,  38  ,  74  , 118  , 142  , 189  ,  87  , 120  , 117  ,  12  , 239  ,
	50  , 106  , 216  , 114  , 137  , 220  ,  42  , 136  , 149  , 158  , 247  , 164  , 190  , 115  , 131  , 128  ,
	20  , 159  , 154  , 249  , 138  , 132  , 109  , 222  , 126  , 127  ,  86  , 150  ,   2  , 185  , 217  , 155  ,
	200  ,  18  ,  98  , 206  , 184  , 101  , 195  , 251  , 224  ,  31  , 113  ,  17  , 188  , 134  ,  61  , 165  ,
	156  , 125  ,  28  , 205  ,  70  ,   3  ,  52  ,  91  ,  13  , 214  , 173  , 198  ,  80  ,  30  , 227  ,   8  ,
	39  ,  83  , 230  ,  27  , 212  , 228  , 163  ,  11  ,   7  , 174  , 234  , 197  , 248  , 144  , 213  , 148  ,
	62  ,  34  ,  23  ,  81  ,  36  , 135  , 202  , 237  , 254  ,  55  , 105  ,  45  ,  37  , 233  ,  19  , 167  ,
	71  , 145  ,  59  ,  51  , 166  ,  54  , 211  ,  22  , 183  , 151  ,  60  ,  16  , 250  , 129  ,   6  , 182  ,
	53  ,  32  ,  25  , 119  ,  57  , 102  , 108  ,  44  , 110  ,  90  , 201  ,  85  ,  68  , 162  ,  21  , 193  ,
	79  ,  41  , 199  , 176  ,  82  , 123  , 246  ,   0  , 180  , 170  ,  95  , 153  ,  64  , 238  , 204  ,  88  ,
};*/


/*byte sbox[256] = {
	1  ,   0  , 140  , 247  , 202  ,  83  , 122  , 208  , 233  ,  78  ,  40  , 193  , 177  , 224  , 228  , 198  ,
	117  , 181  , 171  ,  74  , 152  ,  42  ,  97  ,  94  ,  89  ,  62  , 252  , 205  , 254  ,  65  , 239  , 179  ,
	59  , 111  ,  91  , 240  ,  84  ,  76  , 169  , 200  , 192  ,  11  , 153  ,  20  ,  49  ,  69  , 163  , 195  ,
	45  ,  68  , 147  , 109  , 242  ,  56  , 103  ,  67  , 243  ,  52  ,  33  , 110  , 118  , 186  ,  88  ,  24  ,
	28  , 255  ,  54  , 102  ,  44  ,  48  , 244  , 104  , 166  , 101  , 170  ,  18  ,  85  ,  36  , 232  ,   8  ,
	236  ,  93  ,   4  , 203  ,  77  ,  37  , 134  , 190  ,  25  ,  63  ,  35  , 241  ,  80  , 237  ,  96  ,  22  ,
	23  ,  95  , 174  , 210  ,  72  , 167  ,  55  ,  66  , 245  ,  70  , 144  , 222  ,  50  , 146  ,  32  ,  58  ,
	120  , 182  , 150  , 132  ,  17  , 180  , 187  ,  61  , 183  , 113  , 209  ,   7  , 160  , 251  , 128  , 131  ,
	130  , 127  , 126  , 129  , 151  , 114  , 191  ,  87  , 154  , 159  , 148  , 216  , 246  ,   3  , 184  , 165  ,
	223  , 107  ,  51  , 108  , 217  , 139  , 133  , 115  ,  43  ,  21  , 158  , 137  , 248  , 221  , 136  , 155  ,
	250  , 125  ,  47  , 194  , 142  , 185  , 100  ,  73  ,  39  , 201  ,  19  ,  75  , 207  , 230  , 211  ,  99  ,
	13  , 225  ,  30  , 238  ,  16  , 116  , 121  , 112  , 164  , 143  , 119  ,  60  , 188  , 189  , 135  ,  86  ,
	10  ,  41  ,  46  , 162  , 219  , 213  , 229  ,  14  , 168  ,  38  ,  82  ,   5  ,  26  , 253  , 173  , 231  ,
	123  ,   6  , 175  ,  98  , 196  , 218  , 227  , 235  , 149  , 138  , 197  , 212  , 156  , 249  , 145  , 106  ,
	176  ,  12  , 215  , 234  , 199  ,  15  , 206  , 172  ,   9  ,  79  , 214  , 226  ,  92  ,  81  ,  31  , 178  ,
	90  ,  34  ,  57  ,  53  , 105  ,  71  ,   2  , 141  , 220  , 157  , 124  , 161  , 204  ,  27  ,  64  ,  29  ,
};
*/

byte sbox[sbox_size] = {
};

/*byte sbox[256] = {
	0	
};
*/
byte inv_sbox[sbox_size];

class gen_inv_sbox
{
public:
	gen_inv_sbox()
	{
		for(size_t i = 0; i < sizeof(sbox); ++i)
		{
			inv_sbox[sbox[i]] = i;
		}
	}
};

class gen_random_sbox
{
public:
	gen_random_sbox()
	{
		srand(time(NULL));
		bool* flags = new bool[sizeof(sbox)];
		for(size_t i = 0; i < sizeof(sbox); ++i)
		{
			byte val;
			while(true)
			{
				val = rand() % sizeof(sbox);
				if(!flags[val])
				{
					flags[val] = true;
					break;
				}
			}
			sbox[i] = val;
		}
		delete[] flags;
	}
};

gen_random_sbox g_r_s;

gen_inv_sbox gen;

using namespace std;

extern "C" byte gf_256_11b_mul(byte,byte);

struct sample_ins
{
	byte e;
	vector<byte> or_sbox_out;
	vector<byte> sbox_out;
};

template<class Type>
int get_triangle_rank(Matrix<Type>& m)
{
	int rank = 0;
	for(size_t i = 0; i < m.size1(); ++i)
	{
		bool isZero = true;
		for(size_t j = 0; j < m.size2(); ++j)
		{
			if(!(m(i,j) == 0))
			{
				isZero = false;
				break;
			}
		}
		if(!isZero)
			rank++;
	}
	return rank;
}

int main()
{
	const int eff_size = 4;
	const bool same_eff = true;
	const bool initial_eff = false;
	InitialDimen(8);

	vector<sample_ins> sample_ins_list;

	byte sbox_in[eff_size] = {
		0
	};

	byte sbox_out[eff_size] = {
		0
	};

	byte coeff[eff_size] = {
		1,1,2,1
	};
	if(initial_eff)
	{
		for(int i = 0; i < eff_size; ++i)
		{
			if(same_eff)
			{
				if(i == 0)
				{
					coeff[i] = (rand() % (sbox_size - 1)) + 1;
				}
				else
					coeff[i] = coeff[i-1];
			}
			else
				coeff[i] = i+1;
		}
	}

	srand(time(NULL));
	set<byte> ins_set;

	int new_byte[eff_size] = {
		0
	};

	for(int i = 0; i < eff_size; ++i)
	{
		new_byte[i] = i;
	}

	byte e = rand() % sbox_size; 
	bool full = false;
	
	size_t rank = 0;
	Matrix<SimpleElementS> m(0,sbox_size);
	vector<int> pm,inv_pm;
	pm.resize(sbox_size);
	inv_pm.resize(sbox_size);
	for(size_t i =0 ; i < pm.size(); ++i)
	{
		pm[i] = i;
		inv_pm[i] = i;
	}

	unsigned threhold = 2;

	vector<int> trd_list;
	while(1)
	{
		if(ins_set.size() == (size_t)sbox_size)
		{
			if(!full)
			{
//				printf("Full instance Sample count:%d\n", sample_ins_list.size());
				full = true;
			}
			//	break;
		}
		if( (m.size2() - rank) <= threhold && full)
			break;
/*		if(sample_ins_list.size() >= 40) 
			break;*/

//		byte tmp = sbox[new_byte];

		for(int k = 0; k < 4; ++k)
		{



		sample_ins_list.push_back(sample_ins());
		sample_ins& s_ins = sample_ins_list[sample_ins_list.size() - 1];
		s_ins.e = e;
//		s_ins.or_sbox_out = new_byte;
		m.resize(m.size1(), m.size2()+1);
		pm.resize(pm.size()+1);
		inv_pm.resize(inv_pm.size() + 1);
		pm[pm.size() - 1] = pm.size() - 1;
		inv_pm[inv_pm.size() - 1] = inv_pm.size() - 1;
		for(int i = 0; i < eff_size; ++i)
		{
			s_ins.or_sbox_out.push_back( new_byte[i] );
			byte tmp = sbox[new_byte[i]];
			sbox_in[i] = tmp ^ gf_256_11b_mul(coeff[i], e);
			sbox_out[i] = inv_sbox[sbox_in[i]];
			s_ins.sbox_out.push_back(  sbox_out[i] );

			ins_set.insert(new_byte[i]);
			ins_set.insert(sbox_out[i]);

			m.resize(m.size1()+1, m.size2());
			int equ_pos = m.size1() - 1;
			m(equ_pos, pm[s_ins.or_sbox_out[i]]) = 1;
			m(equ_pos, pm[s_ins.sbox_out[i]]) = 1;
			m(equ_pos, pm[m.size2() - 1]) = coeff[i];

/*			int s1 = rand() % 256;
			int s2 = rand() % 256;
			
			m(equ_pos, pm[s1]) = 1;
			m(equ_pos, pm[s2]) = 1;*/

/*			if(s_ins.or_sbox_out[i] != pm[s_ins.or_sbox_out[i]])
			{
				printf("Caution! not equal!\n");
			}
			if(s_ins.sbox_out[i] != pm[s_ins.sbox_out[i]])
			{
				printf("Caution! not equal!\n");
			}
*/

/*for(size_t row = 0; row < m.size1(); ++row)
			{
				byte val = 0;
				for(size_t col = 0; col < m.size2(); ++col)
				{
					int org_col = inv_pm[col];
					if(org_col < 256)
						val ^= gf_256_11b_mul(m(row, col).getData(), sbox[org_col]);
					else
						val ^= gf_256_11b_mul(m(row,col).getData(), sample_ins_list[org_col-256].e);
				}
				if(val != 0)
				{
					printf("Error copy value !euq_pos:%d;row:%d;sample count:%d;i:%d\n", equ_pos, row, sample_ins_list.size(),i);
					printf("or_sbox_out:%d;sbox_out:%d;pm_oso:%d;pm_so:%d;\n", s_ins.or_sbox_out[i], s_ins.sbox_out[i], pm[s_ins.or_sbox_out[i]], pm[s_ins.sbox_out[i]]);
					printf("Pm:");
					copy(pm.begin(), pm.end(), ostream_iterator<int>(cout, " "));
					printf("\n");
					cout<<m<<endl;
					return 0;
				}
			}*/

			for(size_t row = 0; row < m.size1() - 1; ++row)
			{
				if(row >= rank)
					break;
//				if(m(row,row) == 0)
//					break;
				else
				{
					if(m(equ_pos, row) == 0)
						continue;
					SimpleElementS scale = m(equ_pos, row) / m(row,row);
					for(size_t col = row; col < m.size2(); ++col)
					{
						m(equ_pos, col) -= scale * m(row, col);
					}
				}
			}

			size_t n0_col;
			for(n0_col = 0; n0_col < m.size2(); ++n0_col)
			{
				if(!(m(equ_pos, n0_col) == 0))
				{
					break;
				}
			}
			if(n0_col == m.size2())
			{
				printf("rank equal!");
				m.resize(m.size1() - 1, m.size2());
			}
			else
			{
				SimpleElementS scale = SimpleElementS(1) / m(equ_pos, n0_col);
				for(size_t col = 0; col < m.size2(); ++col)
					m(equ_pos, col) *= scale;
				if(rank != n0_col)
				{
					for(size_t row = 0; row < m.size1(); ++row)
					{
						swap( m(row, rank) , m(row, n0_col));
					}
					swap(inv_pm[rank], inv_pm[n0_col]);
					for(size_t i = 0 ;i < pm.size(); ++i)
					{
						pm[inv_pm[i]] = i;
					}

				}
				for(size_t row = 0; row < rank; ++row)
				{
					if(m(row, rank) == 0)
						continue;
					SimpleElementS scale = m(row, rank) / m(rank, rank);
					for(size_t col = rank; col < m.size2(); ++col)
					{
						m(row, col) -= scale * m(rank, col);
					}
				}

				rank++;
			}


			printf("%d,",rank);

		}
		printf("\n");
//		ins_set.insert(new_byte);


//		printf("%d,",rank);
		fflush(stdout);
		trd_list.push_back(m.size2()-rank);
		printf("trd:%d\n", m.size2() -rank);


//		printf("Sample count:%d,rank:%d,trd:%d\n", sample_ins_list.size(), rank, m.size2() - rank);

		int start = 0;
		for(int i = 0; i < eff_size; ++i)
		{
			for(new_byte[i] = start; new_byte[i] < sbox_size; ++new_byte[i])
			{
				if(ins_set.find(new_byte[i]) == ins_set.end())
					break;
			}
			start = new_byte[i] + 1;
			if(new_byte[i] >= sbox_size)
				new_byte[i] = rand() % sbox_size;
		}
		e = ( rand()) % (sbox_size - 1) + 1;


		}
//		printf("sample size:%d  ins size:%d\n", sample_ins_list.size(), ins_set.size());

	}
	printf("\n\n");
	printf("Sample count:%d\n", sample_ins_list.size());

//	cout<<m<<endl<<endl;

	printf("trd:\n");
	for(size_t i = 0; i < trd_list.size();++i)
		printf("%d,", trd_list[i]);
	printf("\n");

	for(size_t row = 0; row < m.size1(); ++row)
	{
		byte val = 0;
		for(size_t col = 0; col < m.size2(); ++col)
		{
			int org_col = inv_pm[col];
			if(org_col < sbox_size)
				val ^= gf_256_11b_mul(m(row, col).getData(), sbox[org_col]);
			else
				val ^= gf_256_11b_mul(m(row,col).getData(), sample_ins_list[org_col-sbox_size].e);
		}
		if(val != 0)
		{
			printf("Error!equation:%d;sample count:%d\n", row, sample_ins_list.size());
			cout<<m<<endl;
			return 0;
		}
	}

	cout<<m<<endl;

/*	for(size_t i = 0; i < m.size1(); ++i)
	{
		printf("row:%d:\n", i);
		for(size_t j = 0; j < m.size2(); ++j)
		{
			if(!(m(i,j) == 0))
			printf("col:%d,val:%d;", j, m(i,j).getData());
		}
		printf("\n");
	}

//	m.to_triangle();
//	rank = get_triangle_rank(m);
//	printf("rank!!:%d\n",rank);
	return 0;*/

	printf("Sample count:%d, ins count:%d\n", sample_ins_list.size(), ins_set.size());

	Matrix<SimpleElementS> &equ_matrix = m;
//	cout<<equ_matrix<<endl;

	vector<int> free_var_pos;
	for(size_t i = equ_matrix.size1(); i < equ_matrix.size2(); ++i)
	{
		free_var_pos.push_back(i);
		printf("free var row:%d;value:%d\n", i, inv_pm[i]);
	}
	printf("Rank is %d,free var size is %d\n", equ_matrix.size2() - free_var_pos.size(), free_var_pos.size());
	
	int sum = 0;
	rank = equ_matrix.size2() - free_var_pos.size();
	for(size_t i = 0; i < equ_matrix.size1(); ++i)
	{
		for(size_t j = rank; j < equ_matrix.size2();++j)
		{
			if(!(equ_matrix(i,j) == 0))
				++sum;
		}
	}
	printf("free var not zero sum:%d,average:%d\n",sum,sum/rank);
//	cout<<equ_matrix<<endl;

	long long max_try = 1;
	max_try <<= ( free_var_pos.size() * 8); 
	byte *solution = new byte[equ_matrix.size2()];
	for(long long try_val = 0; try_val < max_try; ++try_val)
	{
		memset(solution, 0, equ_matrix.size2());
		if(try_val % 2000 == 0)
			printf("try:%d\n", (int)try_val);
		for(size_t j = 0; j < free_var_pos.size(); ++j)
		{
			int real_col = inv_pm[free_var_pos[j]];
			solution[real_col] = ( try_val >> (8*j) ) & 0xff;
		}
		for(int i = equ_matrix.size1() - 1; i >= 0; --i)
		{
			//free var
			byte result = 0;
/*			for(int j = equ_matrix.size2() - 1; j >= 0 ; --j)
			{
				if(j == i)
					continue;
				result ^= gf_256_11b_mul(equ_matrix(i,j).getData(), solution[j]);
			}
			result = ( SimpleElementS(result) / equ_matrix(i,i) ).getData();*/
			for(size_t j = 0; j < free_var_pos.size(); ++j)
			{
				int col = free_var_pos[j];
				int real_col = inv_pm[free_var_pos[j]];
				result ^= gf_256_11b_mul(equ_matrix(i, col).getData(), solution[real_col]);
			}
			int real_col = inv_pm[i];
			solution[real_col] = result;
		}
		bool equal = true;
		for(int i = 0; i < sbox_size; ++i)
		{
			if(solution[i] != sbox[i])
			{
				equal = false;
				break;
			}
		}
		if(equal)
		{
			printf("Found sbox!\n");
			for(int i = 0; i < 256; ++i)
			{
				printf("%.2x ", solution[i]);
			}
			printf("\n");
			exit(0);
		}
	}
	printf("not found!\n");
//	cout<<equ_matrix<<endl;
	delete[] solution;
	return 0;
}
